# -*- coding: utf-8 -*-
"""
@Time ： 2020-01-06 10:46
@Auth ： chenzj85
@Description：

"""
import shutil

import imageio
import logging
from algo.Algo_interface import Algo_interface
import api as runs
import json
import sys
import os
import glob
from tensorflow import keras
keras.backend.clear_session()
from imageai.Detection.Custom import DetectionModelTrainer
from imageai.Detection.Custom import CustomObjectDetection
import numpy as np
import cv2
import random
import pandas as pd
from detectutils import detectutils
from PIL import Image

class ObjectDetection(Algo_interface):
    def __init__(self, model_type, model_name, model_params):
        self.task_type = model_type
        self.model_name = model_name
        self.model_params = model_params
        self.model = dict()
        self.build_model()
        # return self.model


    def set_model(self, model):
        self.model = model
        return 1

    def get_model(self):
        self.model.pop('train')
        self.model.pop('model')
        return self.model

    def build_model(self):

        if self.model_name == 'yolo':

            self.model['train'] = DetectionModelTrainer()
            self.model['train'].setModelTypeAsYOLOv3()

    def train(self,data):
        def new_report(test_report):
            lists = os.listdir(test_report)  # 列出目录的下所有文件和文件夹保存到lists
            lists.sort(key=lambda fn: os.path.getmtime(test_report + "/" + fn))  # 按时间排序
            file_new = os.path.join(test_report, lists[-1])  # 获取最新的文件保存到file_new
            return file_new

        try:
            xml_train_path,lable_map = data  # 在api.py是tuple进来的两个Series
            self.model['train'].setDataDirectory(data_directory=xml_train_path)
            lable_list = list(lable_map['label'].keys())
            self.model['train'].setTrainConfig(object_names_array=lable_list, batch_size= self.model_params['batch_size'], num_experiments=self.model_params['num_experiments'])
            # In the above,when training for detecting multiple objects,
            # set object_names_array=["object1", "object2", "object3",..."objectz"]
            self.model['train'].trainModel()
            self.model['json'] = os.path.abspath(xml_train_path + '/json/detection_config.json')
            self.model['model_path'] = os.path.abspath(new_report(xml_train_path + '/models'))
            ##训练后的模型加载
            self.model['model'] = CustomObjectDetection()
            self.model['model'].setModelTypeAsYOLOv3()
            self.model['model'].setJsonPath(self.model['json'])
            self.model['model'].setModelPath(self.model['model_path'])
            self.model['model'].loadModel()

        except Exception as e:
            raise Exception(e)

        return 1

    def predict(self,data):   #data为图片路径list
        data,type = data
        if type == 'evaluate':
            def box_detect(image_path):
                detections = self.model['model'].detectObjectsFromImage(input_image=image_path,output_image_path="../detect_image/" + os.path.basename(image_path), minimum_percentage_probability=30)
                predict_box = []
                for eachObject in detections:
                    predict_box.append(eachObject["box_points"])
                return predict_box


            #返回data用于计算出训练端准确率
            data['box_predict'] = list(map(lambda x: box_detect(x), data['image_path']))
            return data

        elif type == 'predict':
            #返回result的dataframe为预测结果可直接写进csv
            result = []
            image_path = list(data['image_path'])
            for image in image_path:
                detections = self.model['model'].detectObjectsFromImage(input_image=image,output_image_path="result.jpg", minimum_percentage_probability=30)
                for eachObject in detections:
                    result.append([os.path.basename(image),eachObject["name"],eachObject["box_points"]])

            return pd.DataFrame(result,columns=['image','label','box_points'])


